{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  import pydrake\n",
    "  import underactuated\n",
    "except ImportError:\n",
    "  !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "  from jupyter_setup import setup_underactuated\n",
    "  setup_underactuated()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trajectory optimization for the Rimless Wheel\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from underactuated.jupyter import SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "from pydrake.examples.rimless_wheel import RimlessWheel\n",
    "from pydrake.all import DirectCollocation, Solve\n",
    "\n",
    "plant = RimlessWheel()\n",
    "context = plant.CreateDefaultContext()\n",
    "\n",
    "params = context.get_numeric_parameter(0)\n",
    "slope = params.slope()\n",
    "alpha = np.pi / params.number_of_spokes()\n",
    "\n",
    "dircol = DirectCollocation(plant,\n",
    "                           context,\n",
    "                           num_time_samples=15,\n",
    "                           minimum_timestep=0.01,\n",
    "                           maximum_timestep=0.1,\n",
    "                           assume_non_continuous_states_are_fixed=True)\n",
    "\n",
    "dircol.AddEqualTimeIntervalsConstraints()\n",
    "\n",
    "dircol.AddConstraintToAllKnotPoints(dircol.state()[0] >= slope - alpha)\n",
    "dircol.AddConstraintToAllKnotPoints(dircol.state()[0] <= slope + alpha)\n",
    "\n",
    "dircol.AddConstraint(dircol.initial_state()[0] == slope - alpha)\n",
    "dircol.AddConstraint(dircol.final_state()[0] == slope + alpha)\n",
    "\n",
    "dircol.AddConstraint(dircol.initial_state()[1] == dircol.final_state()[1] *\n",
    "                     np.cos(2 * alpha))\n",
    "\n",
    "result = Solve(dircol)\n",
    "assert result.is_success()\n",
    "\n",
    "x_trajectory = dircol.ReconstructStateTrajectory(result)\n",
    "\n",
    "x_knots = np.hstack([\n",
    "    x_trajectory.value(t) for t in np.linspace(x_trajectory.start_time(),\n",
    "                                               x_trajectory.end_time(), 100)\n",
    "])\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(x_knots[0, :], x_knots[1, :])\n",
    "ax.plot([x_knots[0, 0], x_knots[0, -1]], [x_knots[1, 0], x_knots[1, -1]], \"--\")\n",
    "\n",
    "# Plot the energy contours.\n",
    "nq = 151\n",
    "nqd = 151\n",
    "mgl = params.mass() * params.gravity() * params.length()\n",
    "q = np.linspace(-0.5, 0.5, nq)\n",
    "qd = np.linspace(-.5, 2, nqd)\n",
    "Q, QD = np.meshgrid(q, qd)\n",
    "Energy = .5 * params.mass() * params.length()**2 * QD**2 + mgl * np.cos(Q)\n",
    "ax.contour(Q,\n",
    "           QD,\n",
    "           Energy,\n",
    "           alpha=0.5,\n",
    "           linestyles=\"dashed\",\n",
    "           colors=\"black\",\n",
    "           linewidths=0.5)\n",
    "\n",
    "ax.set_xlabel(\"theta\")\n",
    "ax.set_ylabel(\"thetadot\")\n",
    "ax.axis([-0.5, 0.5, 0, 2])\n",
    "ax.set_title(\"Limit Cycle of the Rimless Wheel (w/ contours of \"\n",
    "             \"constant energy)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A simple basketball \"bounce pass\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from underactuated.jupyter import SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend('inline')\n",
    "from underactuated import FindResource\n",
    "\n",
    "from pydrake.all import MathematicalProgram, Solve, GetInfeasibleConstraints\n",
    "\n",
    "e = .8    # coefficient of restitution (0 ≤ e ≤ 1)\n",
    "g = 9.8  # gravitational constant (m/s^2)\n",
    "x0 = 0\n",
    "xf = 4\n",
    "z0 = 1\n",
    "zf = 1\n",
    "duration = 5\n",
    "xdot = (xf-x0)/duration\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10,4))\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('z')\n",
    "ax.set_xlim(x0-.1, xf+.1)\n",
    "\n",
    "def bounce_pass(number_of_bounces=0, zdot0_is_positive=True, debug=False):\n",
    "    \n",
    "    prog = MathematicalProgram()\n",
    "    h = prog.NewContinuousVariables(number_of_bounces+1)\n",
    "    prog.AddConstraint(np.sum(h) == duration)\n",
    "    \n",
    "    # zdot at the start of each segment (after collision).\n",
    "    zdot = prog.NewContinuousVariables(number_of_bounces+1)\n",
    "\n",
    "    # Initial velocity constraint (since we *might* have two solutions for each number of bounces)\n",
    "    if (zdot0_is_positive):\n",
    "        prog.AddBoundingBoxConstraint(0, np.inf, zdot[0])\n",
    "    else:\n",
    "        prog.AddBoundingBoxConstraint(-np.inf, 0, zdot[0])\n",
    "    \n",
    "    # Add dynamics constraints for each segment that ends with a bounce\n",
    "    for i in range(number_of_bounces): \n",
    "        # z must be zero at the end of this segment:\n",
    "        z = zdot[i]*h[i] - .5*g*h[i]*h[i] + (z0 if i==0 else 0)\n",
    "        prog.AddConstraint(z == 0)\n",
    "        prog.AddConstraint(zdot[i+1] == -e*(zdot[i] - h[i]*g))\n",
    "        \n",
    "    # Now the final segment\n",
    "    z = zdot[-1]*h[-1] - .5*g*h[-1]*h[-1] + (z0 if number_of_bounces==0 else 0)\n",
    "    prog.AddConstraint(z == zf)\n",
    "    \n",
    "    result = Solve(prog)\n",
    "    if not result.is_success():\n",
    "        if debug:\n",
    "            infeasible = GetInfeasibleConstraints(prog, result)\n",
    "            print(\"Infeasible constraints:\")\n",
    "            for i in range(len(infeasible)):\n",
    "                print(infeasible[i])\n",
    "        # return without plotting\n",
    "        return\n",
    "\n",
    "    # plot the resulting trajectory\n",
    "    ax.set_prop_cycle(plt.rcParams['axes.prop_cycle']) # reset color cycle\n",
    "    relative_time = np.linspace(0, 1, 10)\n",
    "    x_start = x0\n",
    "    for i in range(number_of_bounces+1):\n",
    "        t = result.GetSolution(h[i])*relative_time\n",
    "        x = x_start + t*xdot\n",
    "        z = (z0 if i==0 else 0) + result.GetSolution(zdot[i])*t - .5*g*t*t\n",
    "        ax.plot(x, z, '.-')\n",
    "        x_start = x[-1]    \n",
    "\n",
    "\n",
    "for number_of_bounces in range(0,4):\n",
    "    bounce_pass(number_of_bounces=number_of_bounces, zdot0_is_positive=True)\n",
    "    bounce_pass(number_of_bounces=number_of_bounces, zdot0_is_positive=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The basketball \"trick shot\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from underactuated.jupyter import SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend('inline')\n",
    "from underactuated import FindResource\n",
    "\n",
    "from pydrake.all import MathematicalProgram, Solve, GetInfeasibleConstraints\n",
    "\n",
    "m = 1.\n",
    "r = 0.1\n",
    "g = 9.8  # gravitational constant (m/s^2)\n",
    "e = .8    # coefficient of restitution (0 ≤ e ≤ 1)\n",
    "x0 = -.25\n",
    "xf = -1\n",
    "z0 = 1\n",
    "zf = 3\n",
    "\n",
    "N = 3\n",
    "prog = MathematicalProgram()\n",
    "h = prog.NewContinuousVariables(N)\n",
    "# Make sure time intervals are positive:\n",
    "prog.AddBoundingBoxConstraint(0.01, 10, h)\n",
    "    \n",
    "# zdot at the start of each segment (after collision).\n",
    "xdot = prog.NewContinuousVariables(N)\n",
    "zdot = prog.NewContinuousVariables(N)\n",
    "thetadot = prog.NewContinuousVariables(N)\n",
    "\n",
    "# First segment -- from the launch to the wall.  \n",
    "# Implement the guard (x==0).\n",
    "prog.AddConstraint(x0 + h[0]*xdot[0] == 0).evaluator().set_description('x_wall')\n",
    "z = z0 + h[0]*zdot[0] - 0.5*g*h[0]*h[0]\n",
    "zdot_pre = zdot[0] - g*h[0]\n",
    "# integration + collision dynamics\n",
    "prog.AddConstraint(xdot[1] == -e * xdot[0]).evaluator().set_description('xdot_wall')\n",
    "prog.AddConstraint(zdot[1] == 3.*zdot_pre/5. - 2.*r*thetadot[0]/5.).evaluator().set_description('zdot_wall')\n",
    "prog.AddConstraint(thetadot[1] == -3*zdot_pre/(5.*r) + 2.*thetadot[0]/5.).evaluator().set_description('thetadot_wall')\n",
    "# Help it find the solution we want\n",
    "prog.AddConstraint(zdot[0] >= 0)\n",
    "prog.AddConstraint(thetadot[0] <= 0)\n",
    "\n",
    "# Second segment -- from wall to the floor\n",
    "z = z + zdot[1]*h[1] - 0.5*g*h[1]*h[1]\n",
    "# Implement the guard (z==0).\n",
    "prog.AddConstraint(z == 0).evaluator().set_description('z_floor')\n",
    "x = xdot[1]*h[1]\n",
    "zdot_pre = zdot[1] - g*h[1]\n",
    "# integration + collision dynamics\n",
    "prog.AddConstraint(xdot[2] == 3.*xdot[1]/5. - 2.*r*thetadot[1]/5.).evaluator().set_description('xdot_floor')\n",
    "prog.AddConstraint(zdot[2] == -e*zdot_pre).evaluator().set_description('zdot_floor')\n",
    "prog.AddConstraint(thetadot[2] == -3*xdot[1]/(5.*r) + 2.*thetadot[1]/5.).evaluator().set_description('thetadot_floor')\n",
    "# Make sure the ball doesn't travel too far away\n",
    "prog.AddConstraint(x >= -3.)\n",
    "\n",
    "# Final segment -- from the floor to the hoop\n",
    "x = x + xdot[2]*h[2]\n",
    "z = zdot[2]*h[2] - 0.5*g*h[2]*h[2]\n",
    "# make sure we end at the hoop.\n",
    "prog.AddConstraint(x == xf).evaluator().set_description('x_hoop')\n",
    "prog.AddConstraint(z == zf).evaluator().set_description('z_hoop')\n",
    "zdot_final = zdot[2] - g*h[2]\n",
    "# make sure we're approaching from the correct direction.\n",
    "prog.AddConstraint(xdot[2] >= .1).evaluator().set_description('xdot_hoop')\n",
    "prog.AddConstraint(zdot_final <= -2).evaluator().set_description('zdot_hoop')\n",
    "\n",
    "result = Solve(prog)\n",
    "if not result.is_success():\n",
    "    infeasible = GetInfeasibleConstraints(prog, result)\n",
    "    print(\"Infeasible constraints:\")\n",
    "    for i in range(len(infeasible)):\n",
    "        print(infeasible[i])\n",
    "\n",
    "    # plot the resulting trajectory\n",
    "fig, ax = plt.subplots(figsize=(10,4))\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('z')\n",
    "#ax.set_xlim(x0-.1, xf+.1)\n",
    "ax.set_prop_cycle(plt.rcParams['axes.prop_cycle']) # reset color cycle\n",
    "relative_time = np.linspace(0, 1, 10)\n",
    "x_start = x0\n",
    "z_start = z0\n",
    "for i in range(N):\n",
    "    t = result.GetSolution(h[i])*relative_time\n",
    "    x = x_start + result.GetSolution(xdot[i])*t\n",
    "    z = z_start + result.GetSolution(zdot[i])*t - .5*g*t*t\n",
    "    ax.plot(x, z, '.-')\n",
    "    x_start = x[-1]\n",
    "    z_start = z[-1]\n",
    "    \n",
    "print('Rotation rate during each segment:')\n",
    "print(result.GetSolution(thetadot))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
